import json
import sys
import time

# 读取文件内容
def read_jsonl(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        return [json.loads(line) for line in f]

# 写入文件内容
def write_jsonl(data, file_path):
    with open(file_path, 'w', encoding='utf-8') as f:
        for item in data:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

# 统计 label 分布
def count_label_distribution(data):
    label_count = {}
    for item in data:
        label = item.get("label", "")
        label_count[label] = label_count.get(label, 0) + 1
    return label_count

def main():
    if len(sys.argv) != 4:
        print("用法: python convert_to_standard_format.py <输入文件1>  <输入文件2> <输出文件>", file=sys.stderr)
        sys.exit(1)

    input_file1 = sys.argv[1]
    input_file2 = sys.argv[2]
    output_file = sys.argv[3]

    print("开始读取文件...")
    start_time = time.time()
    # 读取两份 JSONL 文件
    data1 = read_jsonl(input_file1)
    data2 = read_jsonl(input_file2)
    print(f"文件读取完成: {len(data1)} 条来自 {input_file1}, {len(data2)} 条来自 {input_file2}")

    # 构造匹配逻辑
    print("开始匹配数据...")
    match_count = 0
    data2_lookup = {str(d["id"]) + d["conversations"][-1]["value"][:100]: d for d in data2}

    for idx, item in enumerate(data1):
        if idx % 1000 == 0:  # 每 1000 条数据打印一次进度
            print(f"已处理 {idx}/{len(data1)} 条数据...")
        key = str(item["id"]) + item["conversations"][-1]["value"][:100]
        if key in data2_lookup:
            match_count += 1
            data2_lookup[key]["label"] = item.get("label", "")

    print(f"匹配完成，共匹配成功 {match_count} 条数据。")

    # 更新后的数据
    updated_data2 = list(data2_lookup.values())

    # 输出结果
    print("开始写入更新后的文件...")
    write_jsonl(updated_data2, output_file)
    print(f"文件写入完成: {output_file}")

    # 统计 label 分布
    label_distribution = count_label_distribution(updated_data2)
    print("label 分布:")
    for label, count in label_distribution.items():
        print(f"  {label}: {count}")

    end_time = time.time()
    print(f"程序运行完成，耗时 {end_time - start_time:.2f} 秒。")

if __name__ == "__main__":
    main()
